import time
from input_data import *
from model import *
import matplotlib.pyplot as plt
import configparser
import os

os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
config = tf.ConfigProto()
config.gpu_options.per_process_gpu_memory_fraction = 0.8


config.gpu_options.allow_growth = True

def train(image_dir, logs_dir, img_size):
    N_CLASS = 2
    IMG_SIZE = img_size
    BATCH_SIZE = 16
    CAPACITY = 200
    MAX_STEP = 10000
    LEARNING_RATE = 0.0001
    LIST_CHANNELS = [3, 16, 32, 128, 128]

    keep_prob = tf.placeholder(tf.float32)

    # sess = tf.Session(config=config)
    sess = tf.Session()

    train_list = get_train_files(image_dir, True)

    # 计算图定义
    image_train_batch, label_train_batch = get_train_batch(train_list, IMG_SIZE, BATCH_SIZE, CAPACITY, True)
    train_logits = inference(image_train_batch, N_CLASS, LIST_CHANNELS, "train", keep_prob)
    train_loss = losses(train_logits, label_train_batch)
    train_acc = evaluation(train_logits, label_train_batch)
    train_op = tf.train.AdamOptimizer(LEARNING_RATE).minimize(train_loss)

    # 计算参数数目
    var_list = tf.trainable_variables()
    paras_count = tf.reduce_sum([tf.reduce_prod(v.shape) for v in var_list])
    print("参数数目：{}".format(sess.run(paras_count)))

    saver = tf.train.Saver()

    sess.run(tf.global_variables_initializer())

    coord = tf.train.Coordinator()
    threads = tf.train.start_queue_runners(sess=sess, coord=coord)

    fig = plt.figure(figsize=(12, 5))
    ax = fig.add_subplot(1, 2, 1)
    ax.yaxis.grid(True)
    ax.set_title("loss", fontsize=14, y=1.02)
    ax.set_xlabel("step")
    ax.set_ylabel("loss")
    bx = fig.add_subplot(1, 2, 2)
    bx.yaxis.grid(True)
    bx.set_title("accuracy")
    bx.set_xlabel("step")
    bx.set_ylabel("acc")
    step_list = list(range(1001))
    loss_list = []
    acc_list = []

    time_last = time_start = time.time()
    try:
        print("training begins...")
        for step in range(MAX_STEP):
            step = step + 1
            if coord.should_stop():
                break

            _, loss, acc = sess.run([train_op, train_loss, train_acc], feed_dict={keep_prob: 0.4})

            if step % (MAX_STEP / 1000) == 0 or step == 1:
                loss_list.append(loss)
                acc_list.append(acc)

            if step % (MAX_STEP / 100) == 0 or step == 1:
                time_run = time.time() - time_last
                print("    step:%5d,  loss:%.2f,  acc:%.2f%%,  time:%.2fs,  time left:%.1fmin;"
                      % (step, loss, acc * 100, time_run, (MAX_STEP - step) * time_run / 6000))
                time_last = time.time()

            # 每1000步设检查点，保存模型
            if step % (MAX_STEP / 10) == 0 or step == MAX_STEP:
                checkpoint_path = os.path.join(logs_dir, "model.ckpt")
                saver.save(sess, checkpoint_path, global_step=step)

    except tf.errors.OutOfRangeError:
        print("Done.")
    finally:
        coord.request_stop()
    time_end = time.time()
    print("total time:%.1fmin" % ((time_end - time_start) / 60))
    ax.plot(step_list, loss_list)
    bx.plot(step_list, acc_list)
    time_now = list(time.localtime())[0:6]
    figpath = ''.join(['figs\\', '_'.join(str(s) for s in time_now), '.jpg'])
    plt.savefig(figpath)
    plt.show()
    coord.join(threads=threads)
    sess.close()
    print("training over...")


def main():
    config = configparser.ConfigParser()
    config.read('config.ini', encoding='utf-8')
    IMG_SIZE = int(config.get('section_2', 'IMG_SIZE'))
    # DATA = r"G:\Files\PyCharm\Cats_vs_Dogs_2\data\train_"
    DATA = r"data\train"
    LOGS = "logs"
    train(DATA, LOGS, IMG_SIZE)


if __name__ == '__main__':
    main()
